import os
from glob import glob

import cv2
import time
import t_dector_trt

import pycuda.autoinit

import json


def init():
    # Initialize
    # 类别
    t_dector_trt.coco80_names = ['front_wear', 'front_no_wear', 'front_under_nose_wear', 'front_under_mouth_wear',
                               'front_unknown', 'side_wear', 'side_no_wear', 'side_under_nose_wear',
                               'side_under_mouth_wear', 'side_unknown', 'side_back_head_wear', 'side_back_head_no_wear',
                               'back_head', 'mask_front_wear', 'mask_front_under_nose_wear',
                               'mask_front_under_mouth_wear', 'mask_side_wear', 'mask_side_under_nose_wear',
                               'mask_side_under_mouth_wear', 'strap']

    # 加载 模型 /project/train/models/exp/weights/best.onnx
    # engine_file_path =r'/project/train/src_repo/yolov5s.engine'
    engine_file_path = r'/project/train/models/exp/weights/best.engine'

    detector = t_dector_trt.Detector(engine_file_path)  # 加载模型
    model = detector

    return model


def process_image(model, input_image=None, args=None, **kwargs):
    model.cfx.push()  # 3. 推理前执行cfx.push()

    # Do postprocess
    srcimg, list_bboxs = model.detect(input_image)  # 进行 图片 推理

    # cv2.imshow('srcimg_car', srcimg)
    # cv2.waitKey(5)

    model.cfx.pop()

    fake_result = {}

    fake_result["algorithm_data"] = {  # 也就是说，这个这里 只是带上的，不要管
        "is_alert": False,
        "target_count": 0,
        "target_info": []
    }

    fake_result["model_data"] = {"objects": []}

    # Process detections
    cnt = 0
    # 没带口罩的 标签
    no_waers = ['front_no_wear', 'front_under_nose_wear', 'front_under_mouth_wear',
                'side_no_wear', 'side_under_nose_wear', 'side_under_mouth_wear']

    for box in list_bboxs:
        # box = list_bboxs[i]
        left = box[0]
        top = box[1]

        width = box[2]
        height = box[3]
        confidence = float(box[5])
        name = box[4]

        if name in no_waers:  # 没有 带 口罩
            # print(name)
            cnt += 1
            fake_result["model_data"]['objects'].append({  # 这里 才是，真正要添加的
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
            })
            fake_result["algorithm_data"]["is_alert"] = True
            fake_result["algorithm_data"]["target_count"] = cnt

            fake_result["algorithm_data"]["target_info"].append({
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
            })

        else:  # 带了 口罩
            # print(name)
            fake_result["model_data"]['objects'].append({  # 这里 才是，真正要添加的
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
            })

            fake_result["algorithm_data"] = {  # 也就是说，这个这里 只是带上的，不要管
                "is_alert": False,
                "target_count": 0,
                "target_info": []
            }

    return json.dumps(fake_result, indent=4)


if __name__ == '__main__':

    # PLUGIN_LIBRARY = "weights/libmyplugins.so"
    # ctypes.CDLL(PLUGIN_LIBRARY)

    image_names = glob(r'/project/ev_sdk/data/1.jpg')

    # modelpath_car = r'/home/deepin/Documents/ji_pingtai/windows/yolov5/runs/train/exp/weights/best.onnx'
    # image_names = glob(r'/home/deepin/Documents/ji_pingtai/windows/yolov5/data/images/bus.jpg')

    predictor = init()

    s = 0
    count = 0
    for index in range(2):
        for image_name in image_names:
            # print('image_path:', os.path.join(image_dir, image_name))
            img = cv2.imread(image_name)

            start = time.time()

            res = process_image(predictor, img)

            print(res)

            end = time.time()

            inf_end = end - start
            s += end - start

            fps = 1 / inf_end

            fps_label = "FPS: %.2f " % fps
            fps_label = fps_label + " %.2f" % inf_end + "s " + f'({((float(inf_end) * 1000.0)):.1f}ms) Inference'
            print(fps_label)
            print('--------------------------------------------------------:', index)






















